package com.cheetah.start.common.interceptor;

import cn.hutool.core.util.StrUtil;
import com.cheetah.start.common.annotation.RateLimit;
import com.cheetah.start.common.exception.CheetahException;
import com.cheetah.start.common.utils.IpUtil;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;

/**
 * 速率限制 拦截器
 *
 * @author xjh
 * @date 2022/09/05
 */
@Slf4j
@Component
public class RateLimitInterceptor implements HandlerInterceptor {


    /**
     * 不同的接口，不同的流量控制
     * map的key为 Limiter.key
     */
    private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            RateLimit currentLimit = handlerMethod.getMethodAnnotation(RateLimit.class);
            // 存在注解
            if (currentLimit != null) {

                RateLimiter rateLimiter;
                String key;

                String name = currentLimit.name();
                boolean ipRate = currentLimit.ipRate();

                String ipNotHavePort = ":" + IpUtil.getIpNotHavePort(request);
                if (!ipRate) {
                    ipNotHavePort = "";
                }

                if (StrUtil.isNotBlank(name)) {
                    key = name + ipNotHavePort;
                } else {
                    key = request.getContextPath() + ":" + request.getServletPath() + ipNotHavePort;
                }

                if (!limitMap.containsKey(key)) {
                    rateLimiter = RateLimiter.create(currentLimit.permits() - 1);
                    limitMap.put(key, rateLimiter);
                }

                rateLimiter = limitMap.get(key);


                boolean allowTimeOut = currentLimit.allowTimeOut();
                boolean acquire;
                if (allowTimeOut) {
                    acquire = rateLimiter.tryAcquire();
                } else {
                    acquire = rateLimiter.tryAcquire(currentLimit.time(), currentLimit.timeUnit());
                }

                if (!acquire) {
                    throw new CheetahException("请求频繁，请稍后重试！");
                }
            }
        }
        return true;
    }
}
